package mycollections.map;

import mycollections.MyConcurrentModificationException;
import java.util.Comparator;
import java.util.Iterator;

public class MyTreeMap<K, V> implements MyMap<K, V> {
    private final MyEntry NIL=new MyEntry();
    private static final boolean RED = true, BLACK = false;
    private Comparator<K> comparator;
    private int size = 0;
    private MyEntry root;
    private int modificationCount;

    private class MyMapIterator implements Iterator<MyMap.Entry<K,V>> {
        private MyEntry current;
        private int visited = 0;
        private int expectedModificationCount;

        private MyMapIterator(MyEntry currentEntry) {
            this.current = currentEntry;
            expectedModificationCount=MyTreeMap.this.modificationCount;
            clearIsVisited(current);
        }

        private void checkComodification(){
            if (expectedModificationCount != modificationCount)
                throw new MyConcurrentModificationException();
        }

        private void clearIsVisited(MyEntry current) {
            if (current.isNil())
                return;
            current.visited = false;
            clearIsVisited(current.left);
            clearIsVisited(current.right);
        }

        @Override
        public boolean hasNext() {
            checkComodification();
            return visited != size;
        }


        @Override
        public Entry next() {
            checkComodification();
            if (current.left.isNotNil() && !current.left.visited) {
                current = current.left;
                return next();
            }
            if (current.isNotNil() && !current.visited) {
                current.visited = true;
                visited++;
                return current;
            }
            if (current.right.isNotNil() && !current.right.visited) {
                current = current.right;
                return next();
            }
            current = current.parent;
            return next();
        }

        @Override
        public void remove() {
            checkComodification();
            Entry<K,V> oldEntry=next();
            MyTreeMap.this.remove(oldEntry.getKey());
            expectedModificationCount=modificationCount;
        }
    }

    private static class MyEntry<K,V> implements Entry<K, V> {
        private boolean nil;
        private K k;
        private V v;
        private MyEntry<K,V> left, right, parent;
        private boolean visited;
        private boolean red;

        private MyEntry(K k, V v, boolean color, MyEntry<K,V> left, MyEntry<K,V> parent, MyEntry<K,V> right) {
            left.parent=right.parent=this;
            nil =false;
            red = color;
            visited = false;
            this.k = k;
            this.v = v;
            this.left = left;
            this.right = right;
            this.parent = parent;
        }

        private MyEntry(){
            nil =true;
        }

        private MyEntry(MyEntry<K,V> parent){
            this.parent=parent;
            nil =true;
        }

        private boolean isNil() {
            return nil;
        }

        private boolean isNotNil(){
            return !nil;
        }

        private MyEntry grandparent() {
            return parent == null ? null : parent.parent;
        }

        private void setRed() {
            red = true;
        }

        private void setBlack() {
            red = false;
        }

        @Override
        public K getKey() {
            return k;
        }

        @Override
        public V setValue(V value) {
            V oldValue = v;
            v = value;
            return oldValue;
        }

        @Override
        public V getValue() {
            return v;
        }
    }

    private int compareKeys(K k1, K k2) {
        if (comparator != null)
            return comparator.compare(k1, k2);
        else
            return ((Comparable<K>) k1).compareTo(k2);
    }

    public MyTreeMap() {
        comparator = null;
        root=NIL;
        modificationCount=0;
    }

    public MyTreeMap(Comparator<K> comparator) {
        this.comparator = comparator;
        root=NIL;
        modificationCount=0;
    }

    @Override
    public void clear() {
        size = 0;
        root = NIL;
        modificationCount=0;
    }

    @Override
    public boolean containsKey(K key) {
        return get(key) != null;
    }

    @Override
    public boolean containsValue(V value) {
        for (Entry e : this) {
            if (e.getValue().equals(value))
                return true;
        }
        return false;
    }


    @Override
    public V get(K key) {
        MyEntry<K,V> entry= getEntry(key, root);
        return entry!=null?entry.getValue():null;
    }

    private MyEntry<K,V> getEntry(K key, MyEntry<K, V> entry) {
        if (entry.isNil())
            return null;
        int result = compareKeys(key,entry.k);
        if (result == 0)
            return entry;
        else if (result < 0) {
            return getEntry(key, entry.left);
        } else
            return getEntry(key, entry.right);
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public V put(K key, V value) {
        MyEntry<K,V> entry = getEntry(key, root);
        modificationCount++;
        if (entry == null) {
            add(key, value);
            return null;
        } else {
            return entry.setValue(value);
        }
    }

    @Override
    public V remove(K key) {
        MyEntry<K,V> z = getEntry(key, root);
        if (z == null)
            return null;
        V res = z.v;
        remove(z);
        return res;
    }

    private void remove(MyEntry z) {
        MyEntry x, y;
        if (z.left.isNil() || z.right.isNil())
            y = z;
        else y = successor(z);
        if (y.left.isNotNil())
            x = y.left;
        else x = y.right;
        x.parent = y.parent;
        if (y.parent.isNil())
            root = x;
        else if (y == y.parent.left)
            y.parent.left = x;
        else y.parent.right = x;
        if (y != z) {
            z.k = y.k;
            z.v = y.v;
            z.red = y.red;
            z.nil=y.nil;
        }

        if (!y.red)
            removeFixUp(x);
        size--;
        modificationCount++;
    }

    private void removeFixUp(MyEntry x) {
        MyEntry w;
        while (x.isNotNil() && x != root && !x.red) {
            if (x == x.parent.left) {
                w = x.parent.right;
                if (w.red) {
                    w.setBlack();
                    leftRotate(x.parent);
                    w = x.parent.right;
                }
                if (!w.left.red && !w.right.red) {
                    w.setRed();
                    x = x.parent;
                } else {
                    if (!w.right.red) {
                        w.left.setBlack();
                        rightRotate(w);
                        w = x.parent.right;
                    }
                    w.red = x.parent.red;
                    x.parent.setBlack();
                    w.right.setBlack();
                    leftRotate(x.parent);
                    x = root;
                }
            } else {
                w = x.parent.left;
                if (w.red) {
                    w.setBlack();
                    rightRotate(x.parent);
                    w = x.parent.left;
                }
                if (!w.right.red && !w.right.red) {
                    w.setRed();
                    x = x.parent;
                } else {
                    if (!w.left.red) {
                        w.right.setBlack();
                        leftRotate(w);
                        w = x.parent.left;
                    }
                    w.red = x.parent.red;
                    x.parent.setBlack();
                    w.left.setBlack();
                    rightRotate(x.parent);
                    x = root;
                }
            }
        }
        x.setBlack();
    }

    private MyEntry minimum(MyEntry x) {
        while (x.left.isNotNil())
            x = x.left;
        return x;
    }

    private MyEntry successor(MyEntry x) {
        MyEntry y;
        if (x.right.isNotNil())
            return minimum(x.right);
        else {
            y = x.parent;
            while (y.isNotNil() && x == y.right) {
                x = y;
                y = y.parent;
            }
            return y;
        }
    }

    @Override
    public Iterator<MyMap.Entry<K, V>> iterator() {
        return new MyMapIterator(root);
    }

    @Override
    public int size() {
        return size;
    }


    private void add(K k, V v) {
        if (isEmpty()) {
            root = new MyEntry<K,V>(k, v, BLACK, new MyEntry<K,V>(), new MyEntry<K,V>(), new MyEntry<K,V>());
        } else {
            MyEntry<K,V> current = root;
            MyEntry<K,V> prev = current;
            while (current.isNotNil()) {
                prev = current;
                if (compareKeys(k, current.k) < 0)
                    current = current.left;
                else
                    current = current.right;
            }
            MyEntry z = new MyEntry<K,V>(k, v, RED, new MyEntry<K,V>(), prev, new MyEntry<K,V>());
            z.parent = prev;
            if (compareKeys(k, prev.k) < 0)
                prev.left = z;
            else
                prev.right = z;
            insertFixup(z);
        }
        size++;
    }

    private void insertFixup(MyEntry z) {
        MyEntry x, y;
        while (z.parent.isNotNil() && z.grandparent().isNotNil() && z.parent.red) {
            if (z.parent == z.grandparent().left) {
                y = z.grandparent().right;
                if (y.isNotNil() && y.red) {
                    z.parent.setBlack();
                    y.setBlack();
                    z.grandparent().setRed();
                    z = z.grandparent();
                } else {
                    if (z == z.parent.right) {
                        z = z.parent;
                        leftRotate(z);
                    }
                    z.parent.setBlack();
                    z.grandparent().setRed();
                    rightRotate(z.grandparent());
                }
            } else {
                y = z.grandparent().left;
                if (y.isNotNil() && y.red) {
                    z.parent.setBlack();
                    y.setBlack();
                    z.grandparent().setRed();
                    z = z.grandparent();
                } else {
                    if (z == z.parent.left) {
                        z = z.parent;
                        leftRotate(z);
                    }
                    z.parent.setBlack();
                    z.grandparent().setRed();
                    leftRotate(z.grandparent());
                }
            }
        }
        root.setBlack();
    }

    private void leftRotate(MyEntry x) {
        MyEntry<K,V> y = x.right;
        if (root == x)
            root = y;
        else {
            if (x.parent.left == x)
                x.parent.left = y;
            else x.parent.right = y;
        }

        x.right = y.left;
        y.left = x;
        y.parent = x.parent;
        x.parent = y;
    }

    private void rightRotate(MyEntry y) {
        MyEntry x = y.left;
        if (root == y)
            root = x;
        else {
            if (y.parent.left == y)
                y.parent.left = x;
            else y.parent.right = x;
        }

        y.left = x.right;
        x.right = y;
        x.parent = y.parent;
        y.parent = x;
    }
}